#include "order_buffer.h"

#include <stdio.h>
#include <string.h>

#include "util.h"
#ifdef __UNIT_TEST__
#include <cassert>
#endif

Chunk::Chunk() : seq(-1), size(-1), data(nullptr) {}

Chunk::Chunk(const size_t off, const size_t sz, const uint8_t *src) : seq(off), size(sz)
{
    data = static_cast<uint8_t *>(aligned_alloc(ALIGNMENT, alignUp(sz)));
    // 这个Chunk可能不是一个CHUNK_SIZE的倍数
    memcpy(data, src, sz);
}

Chunk::Chunk(const size_t off, const size_t sz) : seq(off), size(sz)
{
    data = static_cast<uint8_t *>(aligned_alloc(ALIGNMENT, alignUp(sz)));
    memset(data, 0, sz);
}

Chunk::Chunk(Chunk &&other) : seq(other.seq), size(other.size), data(other.data)
{
    other.data = nullptr;
}

bool Chunk::operator<(const Chunk &other) { return seq < other.seq; }

Chunk::~Chunk()
{
    free(data);
    data = nullptr;
}

Chunk &Chunk::operator=(Chunk &&other)
{
    if (this != &other)
    {
        free(data);
        seq = other.seq;
        size = other.size;
        data = other.data;
        other.data = nullptr;
    }
    return *this;
}

// 非深拷贝更新(序列号禁止修改)
void Chunk::updateByNonDeepCopy(const size_t sz, uint8_t *&src)
{
    free(data);
    size = sz;
    data = src;
    src = nullptr;
}

void Chunk::clear()
{
    free(data);
    data = nullptr;
}

OrderBuffer::OrderBuffer() {}

OrderBuffer::~OrderBuffer() { clear(); }

void OrderBuffer::addChunk(Chunk &&chunk)
{
    bool mergeRet = false;
    {
        std::lock_guard<std::mutex> lock(mtx);
        // 这里不能随意插入，必须找到合适位置，不然会影响到后续的合并逻辑
        auto it = std::lower_bound(buffer.begin(), buffer.end(), chunk);
        buffer.insert(it, std::move(chunk));
        // 插入时候合并相邻的数据
        mergeRet = mergeAdjacent();
        if (!mergeRet)
        {
            logDebug("current chunks: size: %zu", buffer.size());
        }
    }
    // 精确通知消费者取数据
    consumerCond.notify_all();
}

// 必须按照"期望顺序"处理数据，否则后续写数据到磁盘会出错
void OrderBuffer::getChunk(Chunk &chunk)
{
    const int retryCnts = 3;
    int retryCount = 0;

    std::unique_lock<std::mutex> lock(mtx);

    consumerCond.wait(lock,
                      [this]
                      {
                          return !buffer.empty() && buffer.front().seq * CHUNK_SIZE ==
                                                        expected.load(std::memory_order_relaxed);
                      });
    // 取出数据后，必须更新
    chunk = std::move(buffer.front());
    buffer.pop_front();
    // 能保证下一个期望的Chunk
    expected.fetch_add(chunk.size, std::memory_order_relaxed);
}

// 判断Buffer中是否为空(比如数据被消费完毕)
bool OrderBuffer::isEmpty()
{
    std::lock_guard<std::mutex> lock(mtx);
    return buffer.empty();
}

void OrderBuffer::clear()
{
    std::lock_guard<std::mutex> lock(mtx);
    buffer.clear();
}

// 合并相邻的Chunk，为后续写到磁盘准备
bool OrderBuffer::mergeAdjacent()
{
    // 处理边界(0个、1个Chunk的情况)
    if (buffer.size() < 2)
        return true;

    bool mergeRet = false;
    auto current = buffer.begin();
    while (current != buffer.end())
    {
        auto next = std::next(current);
        // 乱序无法合并
        if (next != buffer.end() &&
            (current->seq * CHUNK_SIZE + current->size == next->seq * CHUNK_SIZE)
            // 最后一个不对齐的一定不要合并，一定要单独写，否则会导致后续无法使用O_DIRECT
            && next->size % CHUNK_SIZE == 0)
        {
            const size_t newSize = current->size + next->size;
            uint8_t *merged = static_cast<uint8_t *>(aligned_alloc(ALIGNMENT, alignUp(newSize)));
            memcpy(merged, current->data, current->size);
            memcpy(merged + current->size, next->data, next->size);
            current->updateByNonDeepCopy(newSize, merged);
            // 成功合并后，释放掉，不要修改current的位置，让其下一轮继续合并
            buffer.erase(next);
            mergeRet = true;
        }
        else  // 遇到不符合合并条件的，向后推进
            current++;
    }

    // 乱序无法合并
    return mergeRet || buffer.size() == 1;
}

void testOrderBuffer()
{
    OrderBuffer orderBuffer;
    {
        logDebug("=========== 测试1 ===========");
        orderBuffer.clear();
        uint8_t dummy_data[CHUNK_SIZE] = {0};
        const int totalSzie = 6 * CHUNK_SIZE - 100;

        orderBuffer.buffer.emplace_back(Chunk(0, CHUNK_SIZE, dummy_data));
        orderBuffer.buffer.emplace_back(Chunk(1, CHUNK_SIZE, dummy_data));
        orderBuffer.buffer.emplace_back(Chunk(3, CHUNK_SIZE, dummy_data));
        orderBuffer.buffer.emplace_back(Chunk(2, CHUNK_SIZE, dummy_data));
        orderBuffer.buffer.emplace_back(Chunk(4, CHUNK_SIZE, dummy_data));
        orderBuffer.buffer.emplace_back(Chunk(5, CHUNK_SIZE - 100, dummy_data));

        logDebug("合并前块数量: %zu", orderBuffer.buffer.size());
        for (const auto &chunk : orderBuffer.buffer)
        {
            logDebug("seq: %zu  , size: %zu", chunk.seq, chunk.size);
        }

#ifdef __UNIT_TEST__
        orderBuffer.mergeAdjacent();
#endif

        logDebug("合并后块数量: %zu", orderBuffer.buffer.size());
        for (const auto &chunk : orderBuffer.buffer)
        {
            logDebug("seq: %zu  , size: %zu", chunk.seq, chunk.size);
        }

        {
            auto it = orderBuffer.buffer.begin();
            auto itNext = std::next(it);
            assert(5 == orderBuffer.buffer.size());
        }

        orderBuffer.clear();
    }

    {
        logDebug("=========== 测试2 ===========");
        orderBuffer.clear();
        uint8_t dummy_data[CHUNK_SIZE] = {0};

        const int totalSzie = 6 * CHUNK_SIZE - 100;

        orderBuffer.buffer.emplace_back(Chunk(0, CHUNK_SIZE, dummy_data));
        orderBuffer.buffer.emplace_back(Chunk(1, CHUNK_SIZE, dummy_data));
        orderBuffer.buffer.emplace_back(Chunk(4, CHUNK_SIZE, dummy_data));
        orderBuffer.buffer.emplace_back(Chunk(5, CHUNK_SIZE - 100, dummy_data));
        orderBuffer.addChunk(Chunk(3, CHUNK_SIZE, dummy_data));
        orderBuffer.addChunk(Chunk(2, CHUNK_SIZE, dummy_data));

        logDebug("合并前块数量: %zu", orderBuffer.buffer.size());
        for (const auto &chunk : orderBuffer.buffer)
        {
            logDebug("seq: %zu  , size: %zu", chunk.seq, chunk.size);
        }

#ifdef __UNIT_TEST__
        orderBuffer.mergeAdjacent();
#endif

        logDebug("合并后块数量: %zu", orderBuffer.buffer.size());
        for (const auto &chunk : orderBuffer.buffer)
        {
            logDebug("seq: %zu  , size: %zu", chunk.seq, chunk.size);
        }

        {
            auto it = orderBuffer.buffer.begin();
            auto itNext = std::next(it);
            assert(2 == orderBuffer.buffer.size() && (totalSzie == it->size + itNext->size) &&
                   (it->size == 5 * CHUNK_SIZE));
        }

        orderBuffer.clear();
    }
}